# 调用ocr 识别 original_img 中的图片， 识别出数据在在9*9 中的坐标
import cv2
import numpy as np
import torch
import torch.nn as nn

# 设置pytesseract的路径（如果需要）
# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

print(torch.cuda.is_available())

class DigitClassifier(nn.Module):
    def __init__(self):
        super(DigitClassifier, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 5 * 5, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 64 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x
# 加载预训练模型
model = DigitClassifier().to(device)
model.load_state_dict(torch.load('digit_ocr_model_pytorch.pth'))
model.eval()


def preprocess_image(image):
    # 转换为灰度图像
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 高斯模糊
    # blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    blurred = gray
    # 二值化处理
    _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    # 反转颜色
    thresh = 255 - thresh
    return thresh

def split_grid(image, rows=9, cols=9):
    height, width = image.shape
    cell_height = height // rows
    cell_width = width // cols
    # 将图像分割成单元格
    print(f"{height} {width} {cell_height} {cell_width}")
    cells = []
    for i in range(rows):
        for j in range(cols):
            # 计算高度范围
            y1 = i * cell_height
            y2 = (i + 1) * cell_height
            # 计算宽度范围
            x1 = j * cell_width
            x2 = (j + 1) * cell_width
            cell = image[y1:y2, x1:x2]

            cell[0:10,:] = 255
            cell[y2-10:, :] = 255
            cell[:, :10] = 255
            cell[:, x2-10:] = 255

            cells.append(cell)
    return cells

def recognize_digit(image_path):
    processed_img = image_path.to(device)
    with torch.no_grad():
        output = model(processed_img)
        probabilities = torch.nn.functional.softmax(output, dim=1)
        confidence, predicted = torch.max(probabilities, 1)


    return predicted.item(), confidence.item()

 # 初始化ORB检测器
# orb = cv2.ORB_create()
def detect_digits(cells):
    digit_positions = []
    i__ = 0
    for i, cell in enumerate(cells):
        min_value = np.min(cell)
        row = i // 9
        col = i % 9
        if min_value< 150:
            i__ = i__+1
            # 调整图像大小以提高
            # 识别精度
            # cell = cv2.resize(cell, (28, 28))


            # print(f"识别结果: {digit_}, 置信度: {confidence:.2f}%")
            cv2.imshow("cell", cell)
            key = cv2.waitKey(0)
            print(f"{key}   {chr(key)}")
            digit_positions.append((row, col, int(chr(key))))
        else:
            digit_positions.append((row, col, int(-1)))

    print(f"{i__ = }")
    return digit_positions

# 读取图像
image = cv2.imread(r"../original_img/001.png")
# 预处理图像
preprocessed = preprocess_image(image)

# cv2.imshow("image", image)
# cv2.imshow("preprocessed", preprocessed)
# cv2.waitKey(0)
# 分割网格
cells = split_grid(preprocessed)

print(len(cells))

# 检测数字
digit_positions = detect_digits(cells)

# 输出数字所在位置
i___ = 0
s=""


for row, col, digit in digit_positions:
    if digit != -1:
        i___ = i___ + 1
        print(f"数字 {digit} 位于第 {row + 1} 行，第 {col + 1} 列")
        s =  s+str(digit)
print(f"87373654136536429394858152")
print(f"{s}")
print(f"{ '87373654136536429394858152' == s }")
print(f"{i___ = }")